using System;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using mshtml;
using OFX.Aggregator.WebScraping.Infrastructure;

namespace OFX.Aggregator.WebScraping.Implementation
{
	internal class HtmlDocumentManager : IDisposable
	{
		private readonly Dictionary<IHTMLDocument, IHtmlDocument> nativeToWrapperDocumentDictionary =
			new Dictionary<IHTMLDocument, IHtmlDocument>();

		private readonly Dictionary<IHTMLElementCollection, IHtmlElementCollection> nativeToWrapperElementCollectionDictionary
			= new Dictionary<IHTMLElementCollection, IHtmlElementCollection>();

		private readonly Dictionary<IHTMLElement, IHtmlElement> nativeToWrapperElementDictionary =
			new Dictionary<IHTMLElement, IHtmlElement>();

		private readonly Dictionary<HTMLWindow2Class, IHtmlFrameWindow> nativeToWrapperFrameDictionary =
			new Dictionary<HTMLWindow2Class, IHtmlFrameWindow>();

		public IHtmlElement GetWrapper(IHTMLElement nativeElement)
		{
			if (nativeElement == null)
				return null;

			IHtmlElement wrapper;
			if (!this.nativeToWrapperElementDictionary.TryGetValue(nativeElement, out wrapper))
			{
				wrapper = new HtmlElementWrapper(nativeElement, this);
				this.nativeToWrapperElementDictionary.Add(nativeElement, wrapper);
			}

			return wrapper;
		}

		public IHtmlDocument GetWrapper(IHTMLDocument nativeDocument)
		{
			if (nativeDocument == null)
				return null;

			IHtmlDocument wrapper;
			if (!this.nativeToWrapperDocumentDictionary.TryGetValue(nativeDocument, out wrapper))
			{
				wrapper = new HtmlDocumentWrapper((HTMLDocumentClass) nativeDocument);
				this.nativeToWrapperDocumentDictionary.Add(nativeDocument, wrapper);
			}

			return wrapper;
		}

		public IHtmlElementCollection GetWrapper(IHTMLElementCollection nativeElementCollection)
		{
			if (nativeElementCollection == null)
				return null;

			IHtmlElementCollection wrapper;
			if (!this.nativeToWrapperElementCollectionDictionary.TryGetValue(nativeElementCollection, out wrapper))
			{
				wrapper = new HtmlElementCollectionWrapper(nativeElementCollection, this);
				this.nativeToWrapperElementCollectionDictionary.Add(nativeElementCollection, wrapper);
			}

			return wrapper;
		}

		public IHtmlFrameWindow GetWrapper(HTMLWindow2Class nativeFrame)
		{
			if (nativeFrame == null)
				return null;

			IHtmlFrameWindow wrapper;
			if (!this.nativeToWrapperFrameDictionary.TryGetValue(nativeFrame, out wrapper))
			{
				wrapper = new HtmlFrameWindowWrapper(nativeFrame, this);
				this.nativeToWrapperFrameDictionary.Add(nativeFrame, wrapper);
			}

			return wrapper;
		}

		public object GetWrapper(object obj)
		{
			if (obj == null)
				return null;

			if (obj is IHTMLElement)
				return GetWrapper((IHTMLElement) obj);
			if (obj is IHTMLElementCollection)
				return GetWrapper((IHTMLElementCollection) obj);
			if (obj is HTMLWindow2Class)
				return GetWrapper((HTMLWindow2Class) obj);
			if (obj is IHTMLDocument)
				return GetWrapper((IHTMLDocument) obj);

			throw new ArgumentException("Unknown HTML object type", "obj");
		}

		#region Implementation of IDisposable

		/// <summary>
		/// Performs application-defined tasks associated with freeing, releasing, or resetting unmanaged resources.
		/// </summary>
		/// <filterpriority>2</filterpriority>
		public void Dispose()
		{
			if (this.nativeToWrapperFrameDictionary != null && this.nativeToWrapperFrameDictionary.Count > 0)
			{
				foreach (HTMLWindow2Class frameWindow in this.nativeToWrapperFrameDictionary.Keys)
				{
					Marshal.ReleaseComObject(frameWindow);
				}

				this.nativeToWrapperElementDictionary.Clear();
			}

			if (this.nativeToWrapperElementDictionary != null && this.nativeToWrapperElementDictionary.Count > 0)
			{
				foreach (IHTMLElement element in this.nativeToWrapperElementDictionary.Keys)
				{
					Marshal.ReleaseComObject(element);
				}

				this.nativeToWrapperElementDictionary.Clear();
			}

			if (this.nativeToWrapperElementCollectionDictionary != null &&
			    this.nativeToWrapperElementCollectionDictionary.Count > 0)
			{
				foreach (IHTMLElementCollection collection in this.nativeToWrapperElementCollectionDictionary.Keys)
				{
					Marshal.ReleaseComObject(collection);
				}

				this.nativeToWrapperElementCollectionDictionary.Clear();
			}

			if (this.nativeToWrapperDocumentDictionary != null && this.nativeToWrapperDocumentDictionary.Count > 0)
			{
				foreach (IHTMLDocument document in this.nativeToWrapperDocumentDictionary.Keys)
				{
					Marshal.ReleaseComObject(document);
				}

				this.nativeToWrapperDocumentDictionary.Clear();
			}
		}

		#endregion
	}
}